import torch
from torch import nn
from torchviz import make_dot
from torch.autograd import Variable
import unittest
from data import *
from utils.augmentations import SSDAugmentation
from layers.modules import MultiBoxLoss
from ssd import build_ssd
import os
import sys
import time
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
import torch.nn.init as init
import torch.utils.data as data
import numpy as np


def make_mlp_and_input():
    model = nn.Sequential()
    model.add_module('W0', nn.Linear(8, 16))
    model.add_module('tanh', nn.Tanh())
    model.add_module('W1', nn.Linear(16, 1))
    x = torch.randn(1, 8)
    return model, x


if __name__ == '__main__':
    """
    这个是绘制网络结构代码，要用cmd打开，并且使用python drawnet.py
    """
    cfg = coco
    model = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
    # model, x = make_mlp_and_input()
    x = torch.randn((1, 3, 300, 300))
    y = model(x)
    dot = make_dot(y, params=dict(model.named_parameters()))
    dot.view()
